"""
Copyright (c) 2024 [XKW.Beijing]
All rights reserved.

Author: [tangxiaojun]
Email: [417281862@qq.com]
"""
from transformers import AutoTokenizer, AutoModelForMaskedLM, BertTokenizer, BertForQuestionAnswering
import torch
import fasttext


class NlpResource:

    def __init__(self):
        self.model = fasttext.load_model('/Users/tangxiaojun/app/src/python-demo/python/nlp/chinese_new_label_model_20240420.bin')
        """
        BERT-Base, Chinese: Chinese Simplified and Traditional, 12-layer, 768-hidden, 12-heads, 110M parameters
        """
        self.tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-chinese")
        self.model = AutoModelForMaskedLM.from_pretrained("google-bert/bert-base-chinese")

        self.tokenizer_qa = BertTokenizer.from_pretrained('ckiplab/bert-base-chinese-qa')
        self.model_qa = BertForQuestionAnswering.from_pretrained('ckiplab/bert-base-chinese-qa')

    def get(self, user_id):
        return {
            "message": f"Hello, World! User ID: {user_id}"
        }

    def predict_model_list(self):
        predictions_text = [
            "游戏行业迎来新一轮监管，加强未成年人保护",
            "奥运会延期举办，运动员备战计划调整",
            "中国一季度GDP上涨5.3%究竟是谁在推动呢？",
            "警方破获一起重大贩毒案件，缴获大量毒品",
            "区块链技术在金融领域应用取得新突破",
            "港股市场迎来反弹，科技股表现抢眼",
            "深圳推出新政策，加强房地产市场调控",
            "政府发布住房保障新措施",
            "股市分析：能源板块迎来投资机会",
            "国产芯片技术取得重大突破",
            "江西国资军工龙头被禁止参与军队采购两年股价一度逼近跌停"
        ]
        result_list = []
        for text in predictions_text:
            # 对文本进行处理，以便于模型预测
            text_split = ' '.join(list(text));
            result_list.append(f"样例 {self.model.predict(text_split)}, {text}")
            # print(f"样例 {self.model.predict(text_split)}, {text}")
        return result_list

    # 预测文本分类，text需要删除标点
    def predict_model_by_text(self, text):
        text_split = ' '.join(list(text));
        return f"样例 {self.model.predict(text_split)}, {text}"

    """
    中文填写
    1、应用 
    2、训练 Hugging Face的transformers库训练自己的语料，微调模型
    例子：https://www.cnblogs.com/zrold/p/17654254.html  https://blog.csdn.net/chehec2010/article/details/136037362
    收集数据集-定义批处理函数，把每句话中指定位置进行挖空（lable,Mark)
    在自然语言处理（NLP）中，特别是在使用像BERT这样的变换器（Transformer）模型时，[MASK]、[SEP]和[CLS]是特殊的标记（tokens），它们在模型的输入处理中扮演着重要的角色。
    """

    def fill_mask(self, input_text):
        # 使用分词器对输入文本进行编码，将 [MASK] 替换为模型的掩码标记
        encoded_input = self.tokenizer(input_text, return_tensors='pt')
        # 获得第一个句子的掩码标记的索引
        mask_token_index = encoded_input['input_ids'].to(torch.device("cuda" if torch.cuda.is_available() else "cpu")).eq(self.tokenizer.mask_token_id)[0]
        # 使用模型进行预测
        with torch.no_grad():
            outputs = self.model(input_ids=encoded_input['input_ids'], attention_mask=encoded_input['attention_mask'])
        # 从模型输出中获取预测结果
        predictions = outputs.logits
        # 将掩码标记的预测结果转换为实际的token id
        predicted_token_id = predictions[0, mask_token_index].argmax(axis=-1)
        # 将token id转换回实际的token
        predicted_token = self.tokenizer.decode(predicted_token_id)
        result = input_text.replace(self.tokenizer.mask_token, predicted_token)
        return result

    """
    提取式问答
    """

    def qa_predict(self, question, context):
        # question = "天安门是哪一年建立的？"
        # context = "天安门（Tian'anmen），坐落在中华人民共和国首都北京市的中心、故宫的南端，与天安门广场以及人民英雄纪念碑、毛主席纪念堂、人民大会堂、中国国家博物馆隔长安街相望，占地面积4800平方米，以杰出的建筑艺术和特殊的政治地位为世人所瞩目。天安门是明清两代北京皇城的正门，始建于明朝永乐十五年（1417年），最初名“承天门”，寓“承天启运、受命于天”之意。设计者为明代御用建筑匠师蒯祥。清朝顺治八年（1651年）更名为天安门。由城台和城楼两部分组成，有汉白玉石的须弥座，总高34.7米。天安门城楼长66米、宽37米。城台下有券门五阙，中间的券门最大，位于北京市皇城中轴线上，过去只有皇帝才可以由此出入。正中门洞上方悬挂着毛泽东画像，两边分别是“中华人民共和国万岁”和“世界人民大团结万岁”的大幅标语。 1925年10月10日，故宫博物院成立，天安门开始对民众开放。1949年10月1日，在这里举行了中华人民共和国开国大典，由此被设计入国徽，并成为中华人民共和国的象征。1961年，中华人民共和国国务院公布为第一批全国重点文物保护单位之一。 [2]"

        # 对问题和上下文进行编码
        encoded_question = self.tokenizer_qa.encode_plus(
            question, context, return_tensors='pt', add_special_tokens=True
        )

        # 禁用dropout等随机性操作
        self.model_qa.eval()

        # 获取模型的输出
        with torch.no_grad():
            outputs = self.model_qa(**encoded_question)

        # 解码模型的输出以获取答案
        answer_start_scores, answer_end_scores = outputs.start_logits, outputs.end_logits
        answer_start = torch.argmax(answer_start_scores)
        answer_end = torch.argmax(answer_end_scores) + 1
        # 提取答案
        answer = self.tokenizer_qa.convert_tokens_to_string(self.tokenizer_qa.convert_ids_to_tokens(encoded_question['input_ids'][0][answer_start:answer_end]))
        return f"问题: {question}\n答案: {answer.replace(' ', '')}"
